import asyncio
import websockets
import json
import socket
from datetime import datetime

BROADCAST_PORT = 12840
enable_captcha = False
captcha_timeout = 30

# 0x01 命令执行
# 0x02 变量操作
# 0x11 存储文件
# 0x12 删除文件
# 0x13 读取文件
# 0x14 列出文件
# 0x15 文件信息
# 0x16 新建文件夹
# 0x90 验证码请求
# 0x91 验证码验证成功
# 0x92 验证码验证失败
# 0x93 验证码验证超时


async def discover_server(timeout=5):
    loop = asyncio.get_running_loop()
    servers = {}

    class DiscoveryProtocol(asyncio.DatagramProtocol):
        def datagram_received(self, data, addr):
            try:
                msg = json.loads(data.decode())
                if msg.get("service") == "afedium_server":
                    server_id = f"{addr[0]}:{msg['port']}"
                    servers[server_id] = {
                        "url": f"ws://{addr[0]}:{msg['port']}",
                        "name": msg.get("server_name", "未知服务器"),
                        "timestamp": datetime.fromisoformat(msg["timestamp"]),
                        "captcha": msg.get("using_captcha", False),
                        "timeout": msg.get("captcha_timeout", 30)
                    }
            except Exception as e:
                print(f"发现服务数据解析失败: {str(e)}")

    transport, protocol = await loop.create_datagram_endpoint(
        lambda: DiscoveryProtocol(),
        family=socket.AF_INET,
        local_addr=('0.0.0.0', BROADCAST_PORT)
    )
    try:
        print(f"正在搜索服务器...")
        await asyncio.sleep(timeout)
    finally:
        transport.close()
        return servers


async def select_server(servers):
    global enable_captcha, captcha_timeout
    if not servers:
        print("未发现任何服务器")
        user_inp = input("请输入服务器URL(address:port)\n: ")
        if input("是否存在验证码(Y/n): ").lower() != 'n':
            enable_captcha = True
        return f"ws://{user_inp}"

    print("发现以下服务器：")
    sorted_servers = sorted(
        servers.values(),
        key=lambda x: x["timestamp"],
        reverse=True
    )

    for idx, server in enumerate(sorted_servers, 1):
        status = "[验证码]" if server["captcha"] else "[开放]"
        print(f"{idx}. {status} {server['name']} ({server['url']})")
    print(f"{len(sorted_servers)+1}. 手动连接")

    while True:
        try:
            choice = await asyncio.get_event_loop().run_in_executor(
                None, input, "请输入要连接的服务器编号 (0取消): "
            )
            choice = int(choice)
            if choice == 0:
                return None
            elif 1 <= choice <= len(sorted_servers):
                selected = sorted_servers[choice-1]
                enable_captcha = selected["captcha"]
                captcha_timeout = selected["timeout"]
                return selected["url"]
            elif choice == 3:
                user_inp = input("请输入服务器URL(address:port)\n: ")
                if input("是否存在验证码(Y/n): ").lower() != 'n':
                    enable_captcha = True
                return f"ws://{user_inp}"
            print("输入无效，请重新输入")
        except ValueError:
            print("请输入有效数字")


async def chat_client():
    servers = await discover_server()
    server_url = await select_server(servers)
    if not server_url:
        print("已取消连接")
        return

    reconnect_delay = 1
    max_reconnect_attempts = 5
    attempt = 0
    while attempt < max_reconnect_attempts:
        try:
            print(f"尝试连接服务器({attempt + 1}/{max_reconnect_attempts})...")
            async with (websockets.connect(
                    server_url,
                    proxy=None,
                    ping_interval=30,
                    ping_timeout=5,
                    close_timeout=5
            ) as websocket):
                print("连接成功!")

                if enable_captcha:
                    # 接收验证请求
                    verify_msg = await websocket.recv()
                    if verify_msg == chr(0x90):
                        print(f"请从服务器控制台获取验证码并输入")

                        try:
                            # 带超时的输入
                            user_code = await asyncio.wait_for(
                                asyncio.get_event_loop().run_in_executor(
                                    None, input, "验证码 >> "
                                ),
                                timeout=captcha_timeout
                            )
                            await websocket.send(chr(0x90) + user_code.strip())

                            # 获取验证结果
                            auth_result = await websocket.recv()
                            if auth_result == chr(0x91):
                                print("认证成功!")
                            elif auth_result == chr(0x91):
                                print("认证失败，连接终止")
                                return
                            elif auth_result == chr(0x93):
                                print("验证码输入超时")
                                return
                            else:
                                print("未知错误")
                                return
                        except asyncio.TimeoutError:
                            print("验证码输入超时")
                            return

                attempt = 0
                while True:
                    try:
                        message = await asyncio.get_event_loop().run_in_executor(
                            None, input, "$: "
                        )
                        if message.replace(' ', '') == '':
                            continue
                        await asyncio.wait_for(
                            websocket.send(chr(eval(message[0:4])) + message[5:]),
                            timeout=10
                        )
                        try:
                            response = await asyncio.wait_for(
                                websocket.recv(),
                                timeout=30
                            )
                            print(response, end='')
                        except asyncio.TimeoutError:
                            print("响应超时，服务器未在30秒内回复")
                    except websockets.exceptions.ConnectionClosed as e:
                        print(f"\n连接关闭: {e.code} - {e.reason}")
                        break
                    except KeyboardInterrupt:
                        print("\n用户中断，正在关闭连接...")
                        await websocket.close()
                        return
                    except Exception as e:
                        print(f"\n发生错误: {str(e)}")
                        break
        except (ConnectionRefusedError, OSError):
            print("无法连接到服务器")
        except websockets.exceptions.InvalidURI:
            print("无效的服务器地址")
            return
        except Exception as e:
            print(f"连接错误: {str(e)}")
        attempt += 1
        if attempt < max_reconnect_attempts:
            print("尝试重新连接...")
            await asyncio.sleep(reconnect_delay)
    print("达到最大重连次数，终止程序")


if __name__ == "__main__":
    try:
        asyncio.run(chat_client())
    except KeyboardInterrupt:
        print("\n客户端已关闭")
    except Exception as e:
        print(f"客户端异常: {str(e)}")
